import os

import numpy as np
from sklearn.metrics import roc_curve, auc
from scipy.interpolate import interp1d
from itertools import cycle
from sklearn.neural_network import MLPClassifier

import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import AdaBoostClassifier
from sklearn.metrics import accuracy_score
from sklearn.svm import SVC

import Index.dataset
from Index.dataset import DataSet
"""
运行adaboost文件
当特征值过多时，运算量会很大，耗时很久
"""
if __name__ == '__main__':

    df = DataSet(ts_code = '000001.SZ', start_date = '20200402', end_date = '20230102').data
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df['timstamp'] = df['trade_date'].astype('int64') / 10 ** 9
    df.drop('trade_date', axis = 1, inplace = True)
    df = df.iloc[:, 0:30]
    X = df.drop(columns='Fmark')
    y = df['Fmark']
    y = y.astype('int64')
    X = X[0:-1]
    y = y[1:]
    #print(y)

    """
    将多个股票叠起来做分类预测
    """
    #X, y = Index.dataset.StackedDataSet(['000001.SZ', '000006.SZ'], '20200201', '20210201')

    Xtrain, Xtest, ytrain, ytest = train_test_split(X, y, test_size=0.2,shuffle=False)

    svc = SVC(probability=True, kernel='linear')

    # 决策树
    #clf = AdaBoostClassifier(random_state=123)

    # csv
    clf = AdaBoostClassifier(estimator=svc, n_estimators=10, learning_rate=1)
    print('training')
    clf.fit(Xtrain, ytrain)
    print('done')
    ypred = clf.predict(Xtest)
    a = pd.DataFrame()
    a['pre'] = list(ypred)
    a['real'] = list(ytest)

    score = accuracy_score(ypred, ytest)
    ypred_proba = clf.predict_proba(Xtest) # [-2. -1.  1.  2.]
    n_classes = ypred_proba.shape[1]
    fpr = dict()
    tpr = dict()
    roc_auc = dict()
    Cclass_name = ['下降', '底', '顶', '上升']
    class_name = [-2, -1, 1, 2]
    for i in range(len(class_name)):
        fpr[i], tpr[i], _ = roc_curve(ytest == class_name[i], ypred_proba[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])

    """
    多特征值画出ROC曲线的两种方法，macro与micro
    """
    # 方法1

    all_fpr = np.unique(np.concatenate([fpr[i] for i in range(n_classes)]))
    # Then interpolate all ROC curves at this points
    mean_tpr = np.zeros_like(all_fpr)
    for i in range(n_classes):
        interpolator = interp1d(fpr[i], tpr[i])
        mean_tpr += interpolator(all_fpr)
    mean_tpr /= n_classes
    fpr["macro"] = all_fpr
    tpr["macro"] = mean_tpr
    roc_auc["macro"] = auc(fpr["macro"], tpr["macro"])

    # 方法2
    unique_classes = np.unique(ytest)
    class_to_int = {cls: i for i, cls in enumerate(unique_classes)}
    int_to_class = {i: cls for cls, i in class_to_int.items()}
    ytest_int = np.array([class_to_int[cls] for cls in ytest])
    num_classes = len(unique_classes)
    L = np.eye(num_classes)[ytest_int]
    # print(L)
    fpr["micro"], tpr["micro"], _ = roc_curve(L.ravel(), ypred_proba.ravel())
    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])

    lw = 2

    plt.figure()
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
    plt.rcParams['axes.unicode_minus'] = False   # 正常显示负号
    plt.plot(fpr["micro"], tpr["micro"],
             label='micro-average ROC curve (area = {0:0.2f})'
                   ''.format(roc_auc["micro"]),
             color='deeppink', linestyle=':', linewidth=4)

    plt.plot(fpr["macro"], tpr["macro"],
             label='macro-average ROC curve (area = {0:0.2f})'
                   ''.format(roc_auc["macro"]),
             color='navy', linestyle=':', linewidth=4)

    colors = cycle(['aqua', 'darkorange', 'cornflowerblue', 'red'])
    for i, color in zip(range(n_classes), colors):
        plt.plot(fpr[i], tpr[i], color=color, lw=lw,
                 label='ROC curve of class {0} (area = {1:0.2f})'
                       ''.format(Cclass_name[i], roc_auc[i]))

    plt.plot([0, 1], [0, 1], 'k--', lw=lw)
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('000002.SZ by csvBase')
    plt.legend(loc="lower right")
    plt.show()